/**
 * @file batchMem.hpp
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-11-09
 *
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 * This file can not be copied or distributed without the permission of Lynxi
 * Technologies Co., Ltd.
 *
 */

#pragma once

#include <lyn_api.h>
#include <malloc.h>
#include <semaphore.h>
#include <cassert>
#include <condition_variable>
#include <list>
#include <mutex>
#include <queue>
#include "logging.h"
#include "util.hpp"

// 批量内存对象，避免批量操作申请多次内存
/********************************************************************************
类名 : BatchMem
Description: BatchMem用来进行批量内存的管理
Example: BatchMem oMemory(false, 10, 5)
*******************************************************************************/
class BatchMem {
 public:
  /**
   * @brief 构造函数，根据类型来决定申请主机侧还是设备侧内存
   *
   * @param[in] bHostType true表示是host内存, false表示是设备侧内存
   * @param[in] elementSize 表示单个内存元素大小
   * @param[in] elementCount 表示内存元素个数
   * @return 无
   */
  BatchMem(bool hostType, int elementSize, int elementCount)
      : m_bHostType(hostType),
        m_iElementSize(elementSize),
        m_iElementCount(elementCount),
        m_pBuffer(nullptr),
        m_iElemUsedCout(0) {
    if ((elementSize > 0 && elementCount > 0) == false) {
      loge("elementSize and elementCount must greater then 0.");
      quick_exit(-1);
    }
    m_iBufferSize = m_iElementSize * m_iElementCount;
    if (m_bHostType) {
      for (int i = 0; i < m_iElementCount; i++) {
        m_vBuffer.emplace_back(malloc(m_iElementSize));
      }
    } else {
      CHECK_ERR(lynMalloc((void **)&m_pBuffer, m_iBufferSize));
      for (int i = 0; i < m_iElementCount; i++) {
        m_vBuffer.push_back((char *)m_pBuffer + m_iElementSize * i);
      }
    }
  }

  /**
   * @brief 析构函数，进行分配内存的释放
   *
   * @param 无
   * @return 无
   */
  ~BatchMem() {
    if (m_bHostType) {
      for (auto i : m_vBuffer) {
        free(i);
      }
    } else {
      if (m_pBuffer) {
        if (lynFree(m_pBuffer) != 0) {
          loge("BatchMem lynFree error");
        }
      }
    }

    m_pBuffer = nullptr;
    m_vBuffer.clear();
  }

  //
  void Reset() {
    std::lock_guard<std::mutex> lock(m_oMutex);
    m_iElemUsedCout = 0;
  }

  // 获取buffer首地址
  void *Buffer() { return m_pBuffer; }

  void *Buffer(int index) {
    std::lock_guard<std::mutex> lock(m_oMutex);
    return (index <= m_iElemUsedCout) ? m_vBuffer[index] : nullptr;
  }

  // 获取整个buffer大小
  int BufferSize() { return m_iBufferSize; }

  // 获取单个元素大小
  int ElementSize() { return m_iElementSize; }

  // 获取元素总数
  int ElementCount() { return m_iElementCount; }

  // 获取已使用的元素
  int Size() {
    std::lock_guard<std::mutex> lock(m_oMutex);
    return m_iElemUsedCout;
  }

  // 获取一个元素
  void *GetElement() {
    std::lock_guard<std::mutex> lock(m_oMutex);
    if (m_iElemUsedCout >= m_iElementCount) {
      return nullptr;
    }

    return m_vBuffer[m_iElemUsedCout++];
  }

 private:
  std::mutex m_oMutex;
  bool m_bHostType;
  int m_iElemUsedCout;            // 当前已用元素个数
  std::vector<void *> m_vBuffer;  // 所有单个内存元素地址
  void *m_pBuffer;                // 内存首地址
  int m_iBufferSize;              // 内存总大小
  int m_iElementSize;             // 单个内存元素大小
  int m_iElementCount;            // 内存元素总数
};                                // class BatchMem

// 批量内存对象池类，避免反复申请批量内存对象
/********************************************************************************
类名 : BatchMemPool
Description:
BatchMemPool就是批量内存对象池端管理类，bHostType表示是host内存还是设备侧内存，
 elementSize表示单个内存元素大小，elementCount表示内存元素个数，iPoolSize表示内存池大小
Example: BatchMemPool oMemory(false, 10, 5, 10)
*******************************************************************************/
class BatchMemPool {
 public:
  // bHostType 为批量内存类型，true表示host端内存，false表示device端内存
  BatchMemPool(bool bHostType, int iElemSize, int iElemCount,
               int iPoolSize = 9999)
      : m_bHostType(bHostType),
        m_iElemSize(iElemSize),
        m_iElemCount(iElemCount),
        m_iCurPoolSize(0),
        m_iPoolSize(iPoolSize) {}

  ~BatchMemPool() { Release(); }

  /**
   * @brief Get the Batch Mem object
   *
   * @return BatchMem*
   */
  BatchMem *GetBatchMem() {
    std::unique_lock<std::mutex> lock(m_oMutex);

    while (m_listBufPool.size() == 0 && m_iCurPoolSize == m_iPoolSize)
      m_oCondVar.wait(lock);

    if (m_listBufPool.size() == 0) {
      BatchMem *pBuf = new BatchMem(m_bHostType, m_iElemSize, m_iElemCount);
      ++m_iCurPoolSize;
      return pBuf;
    } else {
      BatchMem *pBuf = m_listBufPool.front();
      m_listBufPool.pop_front();
      return pBuf;
    }
  }

  /**
   * @brief put the Batch Mem object
   *
   * @param  pData
   */
  void PutBatchMem(BatchMem *pData) {
    std::unique_lock<std::mutex> lock(m_oMutex);
    pData->Reset();
    m_listBufPool.push_back(pData);
    m_oCondVar.notify_one();
  }

  /**
   * @brief Release memory
   *
   */
  void Release() {
    for (BatchMem *pData : m_listBufPool) {
      delete (pData);
    }
    m_listBufPool.clear();
  }

  /**
   * @brief Get Pool Size
   *
   * @return int Pool Size
   */
  int Size() {
    std::lock_guard<std::mutex> lock(m_oMutex);
    return m_listBufPool.size();
  }

 protected:
  std::mutex m_oMutex;
  std::condition_variable m_oCondVar;
  bool m_bHostType;                     //  内存池模式
  int m_iPoolSize;                      // 内存池大小
  int m_iCurPoolSize;                   // 当前内存池大小
  int m_iElemSize;                      // 元素大小
  int m_iElemCount;                     // 元素个数
  std::list<BatchMem *> m_listBufPool;  // 内存池对象
};
